import os
import glob
from random import randint as ri
import cv

false_matches_filename = "bg.txt"
elements = ["slider", "knob", "pad"]
nb_negative_samples = 500
nb_positive_samples = 1000


def generateSamples(element):
    """ generate positive samples for the 'element' in a sub folder """
    os.chdir(element)
    img = cv.LoadImage(element+".png")
    w, h = img.width, img.height
    os.system("""opencv_createsamples -img %s.png -bgcolor 255 -vec %s_positives.vec -maxzangle 0.2 -maxxangle 0.2 -maxyangle 0.2 -num %d -w %d -h %d"""%(element, element, nb_positive_samples, w, h))
    os.chdir("..")

def generateCascade(element):
    """ generate the xml cascade for the 'element' in a sub folder """
    os.chdir(element)
    img = cv.LoadImage(element+".png")
    w, h = img.width, img.height
    os.system("""opencv_haartraining -data %s_cascade -vec %s_positives.vec -bg ..\\bg.txt -npos %d -nneg %d -mode ALL -nonsym -w %d -h %d -nstages 20 -nsplits 2 -minhitrate 0.999 -maxfalsealarm 0.5"""%(element, element,
                                                                                                                 nb_positive_samples,
                                                                                                                 nb_negative_samples,
                                                                                                                 w, h))
    os.chdir("..")

def generateRandomNegatives(element):
    """ generates nb negatives images randomly (so hopefully negative...) of size ~4x the original element """
    imageElem = cv.LoadImage(element+"/"+element+".png")
    iw = imageElem.width
    ih = imageElem.height
    #removing previous images
    map(os.remove, glob.glob("negative_samples"+"/"+"neg*"))
    for n in xrange(nb_negative_samples):
        w = ri(iw, iw*3)
        h = ri(ih, ih*3)
        img = cv.CreateImage((w, h), 8, 1)
        cv.Rectangle(img, (0, 0), (img.width, img.height), 255, -1)
        for lineNb in xrange(ri(1, 10)):
            x1, y1, x2, y2, t = [ri(0, w), ri(0, h), ri(0, w), ri(0, h), ri(1, 3)]
            cv.Line(img, (x1, y1), (x2, y2), 0, t, cv.CV_AA)
        cv.SaveImage("negative_samples/"+"neg%d.png"%n, img)
    
    #write the name of all the files in the negative_samples dir into bg.txt
    with open(false_matches_filename, "w") as f:
        filePaths = ["".join([os.path.join("negative_samples", path), "\n"]) for path in os.listdir("negative_samples") if path!=".svn"]
        f.writelines(filePaths)

if __name__=="__main__":
    for element in elements:
        print "generating negative samples for element %s"%element
        generateRandomNegatives(element)
        print "generating positive samples for element %s"%element
        generateSamples(element)
        print "generating cascade for %s, come back in two days..."%element
        generateCascade(element)
        print "wow it actually finished the %s cascade!"%element
        